from diffusers import StableDiffusionInpaintPipeline
import torch
from PIL import Image,ImageDraw
import numpy as np
import random
import os
from inpainting_dataset import generate_mask_for_outpainting,random_irregular_mask_image

scale_index= 1

def generate_mask(x1y1x2y2):
        x1 = x1y1x2y2[0]
        y1 = x1y1x2y2[1]
        x2 = x1y1x2y2[2]
        y2 = x1y1x2y2[3]
    
        mask = np.zeros((360*scale_index,640*scale_index),dtype=np.uint8)
        mask[int(y1):int(y2),int(x1):int(x2)] = 255
        return Image.fromarray(mask).convert('L')


checkpoint_iters = [

    [300500],
    ]
model_dirs = [
    'sd_inpainting/job_for_safeseat', 
    ]

for iters ,model_dir in zip(checkpoint_iters,model_dirs):
    for ir in iters:
    
        cmd = f'cp /mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting/{model_dir}/checkpoint-{ir}/model.safetensors /mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting/weights/unet/diffusion_pytorch_model.safetensors'
        print(cmd)
        os.system(cmd)
        pipe = StableDiffusionInpaintPipeline.from_pretrained(
            "/mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting/weights",
            torch_dtype=torch.float16,requires_safety_checker=False
        )
        pipe.safety_checker=None
        pipe.to("cuda")
        
        # image = Image.open('/mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting/bg845_726_1115_902.jpg').convert('RGB')
        # image = Image.open('/mnt/afs2d01/luotianhang/diffusion_data/process/data_inpainting/inpainting_material/empty_cart_sort/vc1/001_A/frame_45.jpg').convert('RGB')
        image = Image.open('/mnt/afs2d01/luotianhang/diffusion_data/process/data_inpainting/inpainting_material/empty_cart_sort/a19/empty_1/frame_15.jpg').convert('RGB')
        mem = {}
        pad= 100
        offset_x=-300
        offset_y=-100

        mem['Forward-Facing SafeSeat']=[757-0-offset_x, 302-0-offset_y, 1134+0-offset_x, 1080+0-offset_y]
        mem['Forward-Facing safeset']=[757-0-offset_x, 302-0-offset_y, 1134+0-offset_x, 1080+0-offset_y]
        mem['Rear-Facing SafeSeat']=[757-0-offset_x, 302-0-offset_y, 1134+0-offset_x, 1080+0-offset_y]
        mem['Rear-Facing safeset']=[757-0-offset_x, 302-0-offset_y, 1134+0-offset_x, 1080+0-offset_y]
        mem['a SafeSeat']=[757-0-offset_x, 302-0-offset_y, 1134+0-offset_x, 1080+0-offset_y]
        mem['a safeset']=[757-0-offset_x, 302-0-offset_y, 1134+0-offset_x, 1080+0-offset_y]
        mem['safeset']=[757-0-offset_x, 302-0-offset_y, 1134+0-offset_x, 1080+0-offset_y]
        mem['an empty safeset']=[757-0-offset_x, 302-0-offset_y, 1134+0-offset_x, 1080+0-offset_y]
        mem['a people sit in a safeseat']=[757-0-offset_x, 302-0, 1134+0-offset_x, 1080+0-offset_y]
        # mem['rabbit']=[845-pad,726-pad,1115+pad,902+pad]
        # mem['bird']=[845-pad,726-pad,1115+pad,902+pad]
        # mem['hedgehog']=[845-pad,726-pad,1115+pad,902+pad]

        # pad= 100
        # mem['aeroplane']=[792-pad, 307-pad, 1135+pad, 750+pad]
        # mem['aeroplane ']=[792-pad, 307-pad, 1135+pad, 750+pad]
        # mem[' aeroplane ']=[792-pad, 307-pad, 1135+pad, 750+pad]
        # mem['airplane']=[792-pad, 307-pad, 1135+pad, 750+pad]
        # mem['airplane ']=[792-pad, 307-pad, 1135+pad, 750+pad]
        # mem[' airplane ']=[792-pad, 307-pad, 1135+pad, 750+pad]
        # mem['a furry dog']=[792, 307, 1135, 750]
        # mem['dog']=[792, 307, 1135, 750]
        # mem['a black tablet']=[752, 718, 895, 926]
        # mem['a laptop with an apple on the screen']=[792, 307, 1135, 750]
        # # mem['computer']=[657, 302, 1080, 754] # 这个词肯能被理解成了集群
        # mem['cell phone ']=[775, 430, 883, 582]
        # mem['purple wallet ']=[775, 430, 883, 582]

        # mem['there is a small dog']=[792, 307, 1175, 750]
        # mem['there is a small white cat ']=[547, 102, 764, 350]
        # # mem['a real pet']=[792, 307, 1175, 750]
        # mem['an aeroplane in the sky']=[547, 102, 764, 350]
        # mem['a black pig with mouth opened']=[792, 307, 1135, 750]

        # mem['a black wallet ']=[775, 430, 883, 582]
        # mem['there is a purse, with money inside ']=[775, 430, 883, 582]
        # mem['a mobile phone ']=[775, 430, 883, 582]
        # mem['a brown cat sitting']=[792, 307, 1135, 750]
        # # mem[' pet ']=[792, 307, 1135, 750]
        # mem['a yellow dog ']=[792, 307, 1135, 750]
        # mem['there is a tablet']=[775, 430, 883, 582]

        # mem['a green bag ']=[547, 102, 764, 350]
        # mem['there is a safe seat with a pillow on it ']=[757-0, 302-0, 1134+0, 1080+0]
        # mem['there is a seat ']=[757-0, 302-0, 1134+0, 1080+0]
        # mem['this is a safeseat ']=[757-0, 302-0, 1134+0, 1080+0]
        # mem['a safe seat ']=[547-0, 102-0, 764+0, 350+0]
        # mem['a safeseat ']=[547-0, 102-0, 764+0, 350+0]
       
        
        # mem['big elephant with a long trunk ']=[792, 307, 1135, 750]
        # mem['a brown bird flying ']=[792, 307, 1135, 750] 
        # mem['clothes  ']=[792, 307, 1135, 750]
        # mem['a bread in a sliver plate ']=[792, 307, 1135, 750]
        # mem[' a birthday cake']=[792, 307, 1135, 750]
       

        width = image.width
        height = image.height
        count = 0
        print('image2image')
        for prompt , loc_x1y1x2y2 in mem.items():
            pad = 0
            x1=loc_x1y1x2y2[0]/width*(640*scale_index)
            y1=loc_x1y1x2y2[1]/height*(360*scale_index)
            x2=loc_x1y1x2y2[2]/width*(640*scale_index)
            y2=loc_x1y1x2y2[3]/height*(360*scale_index)

            x1 = max(x1-pad,0)
            y1 = max(y1-pad,0)
            x2 = min(x2+pad,640*scale_index)
            y2 = min(y2+pad,360*scale_index)

            loc_x1y1x2y2=[int(x1),int(y1),int(x2),int(y2)]
            mask_image = generate_mask(loc_x1y1x2y2)
               
            mask_image.save('./mask_temp.png')
            image = image.resize((640*scale_index,360*scale_index))
            res_image = pipe(
                # prompt='Content harmony, best quality, finely detailed, '+prompt+' ' ,
                # prompt='there is a '+prompt+' in the car' ,
                prompt='realistic,best quality,  finely detailed, photorealistic,'+prompt,
                # prompt=prompt,
                image=image, 
                mask_image=mask_image,
                # guidance_scale=6.0,
                # strength=0.8,
                num_inference_steps=50,
                negative_prompt ="worst quality, low quality, normal quality, bad quality, blurry , ugly, chaos, 2D, cartoon"+'cartoon, illustration, 3d, sepia, painting, cartoons, sketch, (worst quality:2), ((monochrome)), ((grayscale:1.2)), (backlight:1.2), analog, analogphoto,freak,anomaly',
                ).images[0]
            
            res_image=res_image.resize((640*scale_index,360*scale_index))
            # compos = Image.composite(res_image.convert('RGBA'), Image.new('RGBA', image.size,(0,0,0,128)), mask_image)
            # compos = Image.blend(res_image.convert('RGBA'), mask_image.convert('RGBA'), alpha=0.5)
            res_image_draw = ImageDraw.Draw(res_image)
            top_left = (loc_x1y1x2y2[0],loc_x1y1x2y2[1])
            bottom_right = (loc_x1y1x2y2[2],loc_x1y1x2y2[3])
            outline_color = 'green'
            line_width = 2
            res_image_draw.rectangle(
                [top_left,bottom_right],outline=outline_color,width=line_width
            )
            save_dir=os.path.join('/mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting','demo',model_dir,str(ir),'image2image')
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
        
            res_image.save(f"{save_dir}/{str(count)+prompt.replace(' ','_')}.png")
            count+=1
        print('text2image')
        mem['a dog swimming in the pool']=[792, 307, 1135, 750]
        mem['a car in the parking lot']=[792, 307, 1135, 750]
        mem['a eagle is drinking the water']=[792, 307, 1135, 750]
        mem['a horse walking in the grass, and a helicopter is after it'] = [0,0,0,0]
        count= 0
        for prompt  in mem.keys():
        
            prompt = prompt
            mask_image = generate_mask([0,0,image.width,image.height])
               
            
            image = image.resize((640*scale_index,360*scale_index))
           
            res_image = pipe(
                prompt='realistic,best quality,  finely detailed, a photo of '+prompt+'' ,
                # prompt='  '+prompt+'' ,
                image=image, 
                mask_image=mask_image,
                # guidance_scale=5.5,
                strength=0.99,
                num_inference_steps=500,
                negative_prompt ="worst quality, low quality, normal quality, bad quality, blurry , ugly, chaos, 2D, cartoon, disfigured, deformed",

                ).images[0]
            
            res_image=res_image.resize((640*scale_index,360*scale_index))
           
            save_dir=os.path.join('/mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting','demo',model_dir,str(ir),'text2image')
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
        
            res_image.save(f"{save_dir}/{str(count)+prompt}.png")
            count +=1

        os.remove('/mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting/weights/unet/diffusion_pytorch_model.safetensors')